import numpy as np

from core.cuda import cuda_module
from implement.variables.parameter import Parameter
from core.layer import Layer
from utils.common import pair
from utils.functions_collect import conv2d


class Conv2d(Layer):
    def __init__(self, out_channels, kernel_size, stride=1,
                 pad=0, nobias=False, dtype=np.float32, in_channels=None):
        """二维卷积层。

        参数：
            out_channels (int)：输出数组的通道数。
            kernel_size (int 或 (int, int))：卷积核的大小。
            stride (int 或 (int, int))：卷积核的步幅。
            pad (int 或 (int, int))：输入数组的空间填充宽度。
            nobias (bool)：如果为True，则不使用偏置项。
            in_channels (int 或 None)：输入数组的通道数。如果为None，
                参数初始化将推迟到第一次前向数据传递时，届时将确定大小。
        """
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.pad = pad
        self.dtype = dtype

        self.W = Parameter(None, name='W')
        if in_channels is not None:
            self._init_W()

        if nobias:
            self.b = None
        else:
            self.b = Parameter(cuda_module.zeros(out_channels, dtype=dtype), name='b')

    def _init_W(self, xp=np):
        C, OC = self.in_channels, self.out_channels
        KH, KW = pair(self.kernel_size)
        scale = np.sqrt(1 / (C * KH * KW))
        W_data = xp.random.randn(OC, C, KH, KW).astype(self.dtype) * scale
        self.W.data = W_data

    def forward(self, x):
        if self.W.data is None:
            self.in_channels = x.shape[1]
            xp = cuda_module
            self._init_W(xp)

        y = conv2d(x, self.W, self.b, self.stride, self.pad)
        return y

